#include "mainwindow.h"
#include "ip.h"

MainWindow::MainWindow(QWidget *parent) :
    QMainWindow(parent)
{
  createAction();
  createLayout();
  connectSignal();
}

MainWindow::~MainWindow()
{
}

void MainWindow::createLayout()
{
  m_imgLabel = new ImageLabel;
  m_imgLabel->setBackgroundRole(QPalette::Base);
  m_imgLabel->setSizePolicy(QSizePolicy::Ignored, QSizePolicy::Ignored);
  m_imgLabel->setScaledContents(true);
  m_scrollArea = new QScrollArea;
  m_scrollArea->setBackgroundRole(QPalette::Dark);
  m_scrollArea->setWidget(m_imgLabel);

  m_origText = new QTextEdit;

  m_magnifyLabel = new ImageLabel(4);
  m_charImgLabel = new ImageLabel(3);
  m_recognizedLabel = new QLabel;
  m_charTextLabel = new QLabel;
  QFont font;
  font.setPointSize(32);
  font.setBold(true);
  m_charTextLabel->setFont(font);
  m_recognizedLabel->setFont(font);

  QHBoxLayout *textLayout = new QHBoxLayout;
  textLayout->addWidget(m_scrollArea, 3);
  textLayout->addWidget(m_origText,2);

  QHBoxLayout *showLayout = new QHBoxLayout;
  showLayout->addWidget(m_magnifyLabel, 3);
  showLayout->addWidget(new QLabel("Current char in Image"),3);
  QVBoxLayout *imgLayout = new QVBoxLayout;
  imgLayout->addWidget(m_charImgLabel);
  imgLayout->addWidget(m_recognizedLabel);
  showLayout->addLayout(imgLayout, 3);
  showLayout->addWidget(m_charTextLabel,2,Qt::AlignHCenter);
  showLayout->addWidget(new QLabel("Current char in Text"),2);

  QWidget *ctrWidget = new QWidget;
  QVBoxLayout *layout = new QVBoxLayout;
  layout->addLayout(textLayout,7);
  layout->addLayout(showLayout,3);
  ctrWidget->setLayout(layout);

  m_toolbar = addToolBar(tr("File"));
  m_toolbar->addAction(m_loadImageAct);
  m_toolbar->addAction(m_loadTextAct);
  m_toolbar->addAction(m_loadCurrentSet);

  setCentralWidget(ctrWidget);
}

void MainWindow::createAction()
{
  m_loadImageAct = new QAction("Load Image", this);
  m_loadTextAct = new QAction("Load Text", this);
  m_loadCurrentSet = new QAction("Load Training Set", this);
}

void MainWindow::connectSignal()
{
  connect(m_loadImageAct, SIGNAL(triggered()), this, SLOT(askAndLoad()));
  connect(m_loadTextAct, SIGNAL(triggered()), this, SLOT(askAndLoadText()));
  connect(m_loadCurrentSet, SIGNAL(triggered()), this, SLOT(loadCurrentSet()));
  connect(m_imgLabel, SIGNAL(mousePressed()), this, SLOT(showMagnify()));
  connect(m_imgLabel, SIGNAL(mousePressed()), this, SLOT(showCharImgOnPressed()));
  connect(m_origText, SIGNAL(selectionChanged()), this, SLOT(selectionTextChanged()));
  connect(this, SIGNAL(curCharImageChanged()), this, SLOT(showRegcognizedText()));
  connect(this, SIGNAL(imgChanged()), this, SLOT(showImage()));
  connect(this, SIGNAL(magnifyChanged()), this, SLOT(showMagnify()));
}

void MainWindow::askAndLoad()
{
  QString fileName = QFileDialog::getOpenFileName(this, "Choose Image ...");
  if (fileName.isEmpty()) return;

  m_mat = cv::imread(fileName.toStdString());
  QTextStream(stdout) << "type = " << m_mat.type() << " " << CV_8UC3
      << " step = " << m_mat.step << " cols = " << m_mat.cols
      << endl;

  if (m_mat.empty()) {
    QMessageBox::information(this, tr("iOCR"),
                             tr("Cannot load image %1.").arg(fileName));
    return;
  }

  m_tmp = m_mat.clone();
  m_imgLabel->m_point = QPoint(50, 50);

  emit imgChanged();
  emit magnifyChanged();
}

void MainWindow::askAndLoadText()
{
  QString fileName = QFileDialog::getOpenFileName(this, "Choose Text ...");
  if (fileName.isEmpty()) return;
  QFile file(fileName);
  if (!file.open(QIODevice::Text | QIODevice::ReadOnly))
  {
    QMessageBox::information(this, tr("iOCR"),
                             tr("Cannot load text %1.").arg(fileName));
    return;
  }
  QTextStream ts(&file);
  QString text = ts.readAll();
  m_origText->setPlainText(text);
  m_origText->setReadOnly(true);
}

void MainWindow::loadCurrentSet()
{
  m_ocr.loadCurrentSet();
}

void MainWindow::showImage()
{
  m_image = IP::mat2qimage(m_tmp);
  m_imgLabel->setPixmap(QPixmap::fromImage(m_image));
  m_imgLabel->adjustSize();
}

void MainWindow::keyPressEvent(QKeyEvent *ev)
{
  QTextStream(stdout) << "key = " << ev->key() << endl;
  int x = m_imgLabel->m_point.x();
  int y = m_imgLabel->m_point.y();
  switch (ev->key())
  {
  case Qt::Key_G :  // seperate back ground color
    {
      if (m_mat.empty()) break;
      RGB rgb = IP::getRGB(m_mat, y, x);
      QTextStream(stdout) << "RGB = " << rgb.r << " " << rgb.g << " " << rgb.b << endl;
      m_mat = IP::separateGround(m_mat, rgb);
      m_tmp = m_mat.clone();
      emit imgChanged();
      break;
    }
  case Qt::Key_P :  // parse image to get lines and characters
    {
      ParseThread *thread = m_imgDoc.createThread(m_mat);
      connect(thread, SIGNAL(finished()), this, SLOT(parseFinished()));
      thread->start();
      break;
    }
  case Qt::Key_A :  // do both G and P
    {
      RGB rgb = IP::getRGB(m_mat, y, x);
      QTextStream(stdout) << "RGB = " << rgb.r << " " << rgb.g << " " << rgb.b << endl;
      m_mat = IP::separateGround(m_mat, rgb);
      ParseThread *thread = m_imgDoc.createThread(m_mat);
      connect(thread, SIGNAL(finished()), this, SLOT(parseFinished()));
      thread->start();
      break;
    }
  case Qt::Key_I:   // load image
    {
      askAndLoad();
      break;
    }
  case Qt::Key_T:   // load text
    {
      askAndLoadText();
      break;
    }
  case Qt::Key_Z:   // move current character image left
    {
      moveImage(-1);
      break;
    }
  case Qt::Key_X:   // move current character image right
    {
      moveImage(+1);
      break;
    }
  case Qt::Key_N:   // move current character in text left
    {
      moveText(-1);
      break;
    }
  case Qt::Key_M:   // move current character in text right
    {
      moveText(+1);
      break;
    }
  case Qt::Key_K:   // move current character in text & image left
    {
      moveImage(-1);
      moveText(-1);
      break;
    }
  case Qt::Key_L:   // move current character in text & image right
    {
      if (m_imgDoc.atEnd()) break;
      moveImage(+1);
      cv::Mat norm = m_ocr.normalize(m_mat(m_imgDoc.curRect()));
      int d = norm.cols / norm.rows;
      moveText(+d);
      break;
    }
  case Qt::Key_S:   // move current character in text & image right
    {
      m_ocr.save(m_curCharImage, m_curCharText);
      showRegcognizedText();
      break;
    }
  case Qt::Key_U:   // scan matched text & image
    {
      bool match = true;
      cv::Rect r = m_imgDoc.curRect();
      cv::rectangle(m_tmp, r, cv::Scalar(255, 0, 0));
      cv::Mat x;
      QString y, yhat;
      QTextCursor cursor = m_origText->textCursor();
      while (match)
      {
        r = m_imgDoc.curRect();
        x = m_mat(r);
        y = cursor.selectedText();
        yhat = m_ocr.recognize(x);
        match = (y == yhat);
        if (match)
        {
          m_ocr.save(x, y);
          if (m_imgDoc.atEnd()) break;
          m_imgDoc.moveCurrent(1);
          r = m_imgDoc.curRect();
          cv::Mat norm = m_ocr.normalize(m_mat(r));
          int d = norm.cols / norm.rows;
          cursor.clearSelection();
          cursor.movePosition(QTextCursor::NextCharacter, QTextCursor::KeepAnchor,d);
        }
      }
      r = m_imgDoc.curRect();
      cv::rectangle(m_tmp, r, cv::Scalar(0, 0, 255));
      m_curCharImage = m_mat(r);
      m_origText->setTextCursor(cursor);
      emit imgChanged();
      emit curCharImageChanged();
      selectionTextChanged();
      break;
    }
  case Qt::Key_R:   // translate image document
    {
      QString fn = QFileDialog::getSaveFileName(this, "Translate to ...");
      if (fn.isEmpty()) break;
      QFile f(fn);
      if (!f.open(QIODevice::WriteOnly | QIODevice::Text))
        break;
      QTextStream fs(&f);
      fs << "<html>\n"
            "<head>\n"
            "<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\">\n"
            "</head>\n"
            "<body>\n";
      QTextStream(stdout) << "<html>\n"
          "<head>\n"
          "<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\">\n"
          "</head>\n"
          "<body>\n";
      int i = 0;
      Q_FOREACH(const Line& line, m_imgDoc.lines())
      {
        if (i > 0 && line.start - m_imgDoc.lines()[i-1].start > 40)
          fs << "<br/><br/><br/><br/>\n";
        QTextStream(stdout) << "<br/><br/><br/><br/>\n";
        int j = 0;
        Q_FOREACH(const cv::Rect& r, line.chars)
        {
          cv::Mat x = m_mat(r);
          QString y = m_ocr.recognize(x);
          fs << y;
          QTextStream(stdout) <<  y;
          j++;
        }
        fs << "\n";
        QTextStream(stdout) << "\n";
        i++;
      }
      fs << "</body>\n"
            "</html>\n";
      QTextStream(stdout) << "</body>\n"
            "</html>\n";
      break;
    }
  }
}

void MainWindow::showMagnify()
{
  int x = m_imgLabel->m_point.x();
  int y = m_imgLabel->m_point.y();
  int w = 30, w2 = w / 2;
  cv::Mat small(m_mat, cv::Rect(x-w2, y-w2, w, w));
  m_magnifyLabel->show(small);
}

void MainWindow::showLine(int line)
{
  m_tmp = m_mat.clone();
  cv::line(m_tmp, cv::Point(0, line), cv::Point(m_tmp.cols, line), cv::Scalar(0, 255, 0), 2);
  showImage();
}

void MainWindow::lineDetected(int startLine, int endLine)
{
  QTextStream(stdout) << "line = " << startLine << " " << endLine << endl;
}

void MainWindow::parseFinished()
{
  sender()->deleteLater();
  m_tmp = m_mat.clone();
  int mw = 0, mh = 0;
  Q_FOREACH(const Line& line, m_imgDoc.lines())
  {
    if (line.type == 0) continue;
    cv::line(m_tmp, cv::Point(0, line.start), cv::Point(m_tmp.cols, line.start), cv::Scalar(0, 0, 255), 1);
    cv::line(m_tmp, cv::Point(0, line.end), cv::Point(m_tmp.cols, line.end), cv::Scalar(0, 0, 255), 1);
    QTextStream(stdout) << "chars = " << line.chars.size() << endl;
    Q_FOREACH(const cv::Rect& r, line.chars)
    {
      cv::rectangle(m_tmp, r, cv::Scalar(255, 0, 0));
      if (mw < r.width) mw = r.width;
      if (mh < r.height) mh = r.height;
    }
  }
  emit imgChanged();
  QTextStream(stdout) << "mw = " << mw << " mh = " << mh << endl;
}

void MainWindow::moveImage(int d)
{
  if (m_mat.empty()) return;
  cv::Rect r = m_imgDoc.curRect();
  cv::rectangle(m_tmp, r, cv::Scalar(255, 0, 0));
  m_imgDoc.moveCurrent(d);
  r = m_imgDoc.curRect();
  cv::rectangle(m_tmp, r, cv::Scalar(0, 0, 255));
  m_curCharImage = m_mat(r);

  emit imgChanged();
  emit curCharImageChanged();
}

void MainWindow::moveText(int d)
{
  QTextCursor cursor = m_origText->textCursor();
  int p = cursor.position(), a = cursor.anchor();
  cursor.clearSelection();
  if (d > 0)
  {
    if (p < a)
      cursor.movePosition(QTextCursor::NextCharacter, QTextCursor::MoveAnchor, a-p);
    cursor.movePosition(QTextCursor::NextCharacter, QTextCursor::KeepAnchor, d);
  }
  else
  {
    if (p > a)
      cursor.movePosition(QTextCursor::PreviousCharacter, QTextCursor::MoveAnchor, p-a);
    cursor.movePosition(QTextCursor::PreviousCharacter, QTextCursor::KeepAnchor, -d);
  }
  m_origText->setTextCursor(cursor);
  selectionTextChanged();
}

void MainWindow::selectionTextChanged()
{
  QTextCursor cursor = m_origText->textCursor();
  m_curCharText = cursor.selectedText();
  QTextStream(stdout) << m_curCharText << "\n";
  m_charTextLabel->setText(m_curCharText);
  checkRecognizedText();
}

void MainWindow::showCharImgOnPressed()
{
  int i = 0;
  Q_FOREACH(const Line& line, m_imgDoc.lines())
  {
    int j = 0;
    Q_FOREACH(const cv::Rect& r, line.chars)
    {
      if (r.x <= m_imgLabel->m_point.x() && r.x+r.width > m_imgLabel->m_point.x() &&
          r.y <= m_imgLabel->m_point.y() && r.y+r.height > m_imgLabel->m_point.y())
      {
        cv::Rect r = m_imgDoc.curRect();
        cv::rectangle(m_tmp, r, cv::Scalar(255, 0, 0));
        m_imgDoc.setCurLine(i);
        m_imgDoc.setCurChar(j);
        r = m_imgDoc.curRect();
        cv::rectangle(m_tmp, r, cv::Scalar(0, 0, 255));
        showImage();

        m_curCharImage = m_mat(r);
        emit curCharImageChanged();
        return;
      }
      j++;
    }
    i++;
  }
}

void MainWindow::showRegcognizedText()
{
  m_charImgLabel->show(m_ocr.normalize(m_curCharImage));
  QString y = m_ocr.recognize(m_curCharImage);
  y = (y.isEmpty() ? "!!NOT FOUND!!" : y);
  m_recognizedLabel->setText(y);
  QTextStream(stdout) << "recognized = " << y << endl;
  checkRecognizedText();
}

void MainWindow::checkRecognizedText()
{
  QPalette p;
  p.setColor(QPalette::Foreground, Qt::red);
  if (m_recognizedLabel->text() == m_charTextLabel->text())
    p.setColor(QPalette::Foreground, Qt::green);
  m_recognizedLabel->setPalette(p);
  m_charTextLabel->setPalette(p);
}
